"""
标准化流域划分工具
提供分步骤的流域划分流程，输出标准化的GeoJSON格式数据

主要功能：
1. 分步骤生成流程，每个步骤输出明确的中间结果
2. 统一使用GeoJSON格式输出
3. 完整的元数据信息
4. 标准化的文件命名规范

输出文件：
- {outlet_id}_watershed.geojson: 流域边界
- {outlet_id}_rivers.geojson: 河网数据
- {outlet_id}_catchments.geojson: 集水面数据
- outlet_points.geojson: 出口点数据

使用示例:
    from standardized_delineate import StandardizedDelineator
    
    delineator = StandardizedDelineator()
    results = delineator.delineate_watershed(
        lat=64.71072,
        lng=-21.60337,
        outlet_name="Nordhura River at Stekkur",
        outlet_id="6401070",
        area_estimate=507
    )
"""

import json
import os
import sys
import tempfile
import shutil
from datetime import datetime
from typing import Dict, Optional, Any

try:
    import geopandas as gpd
    import pandas as pd
    from loguru import logger
    DEPENDENCIES_AVAILABLE = True
except ImportError:
    DEPENDENCIES_AVAILABLE = False

# 使用相对导入
from ..utils import config
from . import delineate


class StandardizedDelineator:
    """标准化流域划分器"""

    def __init__(self, output_dir: Optional[str] = None):
        """
        初始化标准化流域划分器
        
        Args:
            output_dir: 输出目录，默认使用config中的设置
        """
        self.output_dir = output_dir or config.OUTPUT_DIR
        self.results = {}
        self.processing_steps = []

    def delineate_watershed(self,
                            lat: float,
                            lng: float,
                            outlet_name: str = "Custom_Outlet",
                            outlet_id: str = "1",
                            area_estimate: Optional[float] = None,
                            high_res: bool = True) -> Dict[str, Any]:
        """
        执行标准化流域划分
        
        Args:
            lat: 纬度
            lng: 经度
            outlet_name: 出口点名称
            outlet_id: 出口点ID
            area_estimate: 预估面积（可选）
            high_res: 是否使用高分辨率模式
            
        Returns:
            包含所有输出文件路径和处理结果的字典
        """

        logger.info("=== 开始标准化流域划分流程 ===")

        # 步骤1: 准备输入数据
        self._log_step("准备输入数据和配置")
        input_data = self._prepare_input_data(lat, lng, outlet_name, outlet_id, area_estimate)

        # 步骤2: 执行流域划分
        self._log_step("执行流域划分算法")
        watershed_result = self._execute_delineation(input_data, high_res)

        # 步骤3: 生成标准化输出
        self._log_step("生成标准化GeoJSON输出")
        output_files = self._generate_standardized_outputs(watershed_result, input_data)

        # 步骤4: 创建处理报告
        self._log_step("创建处理报告和元数据")
        report = self._create_processing_report(input_data, output_files)

        # 步骤5: 验证输出
        self._log_step("验证输出文件完整性")
        validation_result = self._validate_outputs(output_files)

        # 步骤6: 复制到最终输出目录（如果配置了）
        final_output_files = output_files.copy()
        if config.FINAL_OUTPUT_DIR and config.FINAL_OUTPUT_DIR.strip():
            self._log_step("复制文件到最终输出目录")
            final_output_files = self._copy_to_final_output(output_files, input_data['outlet_id'])

        logger.success("=== 标准化流域划分流程完成 ===")

        return {
            'input_data': input_data,
            'output_files': output_files,
            'final_output_files': final_output_files,
            'processing_report': report,
            'validation_result': validation_result,
            'processing_steps': self.processing_steps
        }

    def _log_step(self, step_description: str):
        """记录处理步骤"""
        timestamp = datetime.now().isoformat()
        step_info = {
            'timestamp': timestamp,
            'step': len(self.processing_steps) + 1,
            'description': step_description
        }
        self.processing_steps.append(step_info)
        logger.info(f"步骤 {step_info['step']}: {step_description}")

    def _prepare_input_data(self, lat: float, lng: float, outlet_name: str,
                            outlet_id: str, area_estimate: Optional[float]) -> Dict[str, Any]:
        """准备输入数据"""

        # 创建临时目录和CSV文件
        temp_dir = tempfile.mkdtemp()
        temp_csv = os.path.join(temp_dir, 'temp_outlets.csv')

        # 创建出口点数据
        outlet_data = {
            'id': [outlet_id],
            'lat': [lat],
            'lng': [lng],
            'name': [outlet_name]
        }

        if area_estimate:
            outlet_data['area'] = [area_estimate]

        # 保存临时CSV文件
        df = pd.DataFrame(outlet_data)
        df.to_csv(temp_csv, index=False)

        input_data = {
            'outlet_id': outlet_id,
            'outlet_name': outlet_name,
            'coordinates': {'lat': lat, 'lng': lng},
            'area_estimate': area_estimate,
            'temp_dir': temp_dir,
            'temp_csv': temp_csv,
            'created_at': datetime.now().isoformat()
        }

        logger.info(f"输入数据准备完成: {outlet_name} ({lat}, {lng})")
        return input_data

    def _execute_delineation(self, input_data: Dict[str, Any], high_res: bool) -> Dict[str, Any]:
        """执行流域划分"""

        # 保存原始配置
        original_config = self._backup_config()

        try:
            # 设置临时配置
            self._set_temp_config(input_data['temp_csv'], high_res)

            # 重新加载配置和模块
            import importlib
            importlib.reload(config)
            importlib.reload(delineate)

            # 执行流域划分
            logger.info("开始执行流域划分算法...")
            delineate.delineate()

            # 收集结果
            result = {
                'success': True,
                'output_dir': config.OUTPUT_DIR,
                'outlet_id': input_data['outlet_id']
            }

            logger.info("流域划分算法执行完成")
            return result

        except Exception as e:
            logger.error(f"流域划分执行失败: {str(e)}")
            return {'success': False, 'error': str(e)}

        finally:
            # 恢复原始配置
            self._restore_config(original_config)

            # 清理临时文件
            import shutil
            if os.path.exists(input_data['temp_dir']):
                shutil.rmtree(input_data['temp_dir'])

    def _backup_config(self) -> Dict[str, Any]:
        """备份当前配置"""
        return {
            'OUTLETS_CSV': config.OUTLETS_CSV,
            'HIGH_RES': config.HIGH_RES,
            'OUTPUT_DIR': config.OUTPUT_DIR,
            'OUTPUT_EXT': config.OUTPUT_EXT,
            'VERBOSE': config.VERBOSE
        }

    def _restore_config(self, original_config: Dict[str, Any]):
        """恢复原始配置"""
        for key, value in original_config.items():
            setattr(config, key, value)

    def _set_temp_config(self, temp_csv: str, high_res: bool):
        """设置临时配置"""
        os.environ['OUTLETS_CSV'] = temp_csv
        os.environ['HIGH_RES'] = str(high_res)
        os.environ['OUTPUT_DIR'] = self.output_dir
        os.environ['OUTPUT_EXT'] = 'geojson'
        os.environ['VERBOSE'] = 'True'
        os.environ['PLOTS'] = str(config.PLOTS).lower()

    def _generate_standardized_outputs(self, watershed_result: Dict[str, Any],
                                       input_data: Dict[str, Any]) -> Dict[str, str]:
        """生成标准化输出文件"""

        if not watershed_result.get('success'):
            logger.error("流域划分失败，无法生成输出文件")
            return {}

        output_files = {}
        outlet_id = input_data['outlet_id']
        output_dir = watershed_result['output_dir']

        # 查找生成的文件并重命名为标准格式
        if os.path.exists(output_dir):
            for file in os.listdir(output_dir):
                old_path = os.path.join(output_dir, file)

                # 确定文件类型并重命名
                if file.startswith(outlet_id):
                    if '_watershed.geojson' in file:
                        new_name = f"{outlet_id}_watershed.geojson"
                        output_files['watershed'] = os.path.join(output_dir, new_name)
                    elif '_rivers.geojson' in file:
                        new_name = f"{outlet_id}_rivers.geojson"
                        output_files['rivers'] = os.path.join(output_dir, new_name)
                    else:
                        continue

                    # 重命名文件（如果需要）
                    new_path = os.path.join(output_dir, new_name)
                    if old_path != new_path and os.path.exists(old_path):
                        os.rename(old_path, new_path)
                
                # 查找outlet_points.geojson文件（不以outlet_id开头）
                elif file == 'outlet_points.geojson':
                    output_files['outlet_points'] = old_path

        logger.info(f"生成标准化输出文件: {len(output_files)} 个文件")
        return output_files

    def _create_processing_report(self, input_data: Dict[str, Any],
                                  output_files: Dict[str, str]) -> Dict[str, Any]:
        """创建处理报告"""

        report = {
            'processing_info': {
                'tool_name': 'mghydro standardized delineator',
                'version': '2.0.0',
                'processing_date': datetime.now().isoformat(),
                'coordinate_system': 'WGS84 (EPSG:4326)'
            },
            'input_parameters': input_data,
            'output_files': output_files,
            'processing_steps': self.processing_steps,
            'file_descriptions': {
                'watershed': 'Watershed boundary polygon in GeoJSON format',
                'rivers': 'River network within the watershed in GeoJSON format',
                'outlet_points': 'Outlet points with processing results in GeoJSON format'
            }
        }

        # 保存报告到文件
        report_file = os.path.join(self.output_dir, f"{input_data['outlet_id']}_processing_report.json")
        with open(report_file, 'w', encoding='utf-8') as f:
            json.dump(report, f, indent=2, ensure_ascii=False)

        logger.info(f"处理报告已保存: {report_file}")
        return report

    def _validate_outputs(self, output_files: Dict[str, str]) -> Dict[str, Any]:
        """验证输出文件"""

        validation_result = {
            'valid_files': [],
            'invalid_files': [],
            'missing_files': [],
            'file_info': {}
        }

        expected_files = ['watershed', 'outlet_points']

        for file_type in expected_files:
            if file_type in output_files:
                file_path = output_files[file_type]
                if os.path.exists(file_path):
                    try:
                        # 尝试读取GeoJSON文件
                        gdf = gpd.read_file(file_path)
                        validation_result['valid_files'].append(file_type)
                        validation_result['file_info'][file_type] = {
                            'path': file_path,
                            'size_bytes': os.path.getsize(file_path),
                            'feature_count': len(gdf),
                            'crs': str(gdf.crs) if gdf.crs else 'Unknown'
                        }
                        logger.info(f"文件验证通过: {file_type} ({len(gdf)} 个要素)")
                    except Exception as e:
                        validation_result['invalid_files'].append({
                            'file_type': file_type,
                            'path': file_path,
                            'error': str(e)
                        })
                        logger.error(f"文件验证失败: {file_type} - {str(e)}")
                else:
                    validation_result['missing_files'].append(file_type)
                    logger.warning(f"文件不存在: {file_type}")
            else:
                validation_result['missing_files'].append(file_type)

        # 检查可选文件
        if 'rivers' in output_files:
            file_path = output_files['rivers']
            if os.path.exists(file_path):
                try:
                    gdf = gpd.read_file(file_path)
                    validation_result['valid_files'].append('rivers')
                    validation_result['file_info']['rivers'] = {
                        'path': file_path,
                        'size_bytes': os.path.getsize(file_path),
                        'feature_count': len(gdf),
                        'crs': str(gdf.crs) if gdf.crs else 'Unknown'
                    }
                    logger.info(f"河网文件验证通过: {len(gdf)} 个要素")
                except Exception as e:
                    validation_result['invalid_files'].append({
                        'file_type': 'rivers',
                        'path': file_path,
                        'error': str(e)
                    })

        validation_result['is_valid'] = (
                len(validation_result['invalid_files']) == 0 and
                len(validation_result['missing_files']) == 0
        )

        logger.info(f"输出验证完成: {len(validation_result['valid_files'])} 个有效文件")
        return validation_result

    def _copy_to_final_output(self, output_files: Dict[str, str], outlet_id: str) -> Dict[str, str]:
        """复制文件到最终输出目录"""
        
        final_output_files = {}
        final_dir = config.FINAL_OUTPUT_DIR.strip()
        
        if not final_dir:
            logger.warning("FINAL_OUTPUT_DIR 未配置，跳过文件复制")
            return output_files
        
        try:
            # 确保最终输出目录存在
            os.makedirs(final_dir, exist_ok=True)
            logger.info(f"最终输出目录: {final_dir}")
            
            # 复制每个输出文件
            for file_type, source_path in output_files.items():
                if os.path.exists(source_path):
                    filename = os.path.basename(source_path)
                    dest_path = os.path.join(final_dir, filename)
                    
                    # 复制文件
                    shutil.copy2(source_path, dest_path)
                    final_output_files[file_type] = dest_path
                    logger.info(f"已复制 {file_type}: {source_path} -> {dest_path}")
                else:
                    logger.warning(f"源文件不存在，跳过复制: {source_path}")
                    final_output_files[file_type] = source_path
            
            # 同时复制 OUTPUT.csv 文件（如果存在）
            csv_source = os.path.join(self.output_dir, "OUTPUT.csv")
            if os.path.exists(csv_source):
                csv_dest = os.path.join(final_dir, f"{outlet_id}_OUTPUT.csv")
                shutil.copy2(csv_source, csv_dest)
                final_output_files['csv_summary'] = csv_dest
                logger.info(f"已复制 CSV 摘要: {csv_source} -> {csv_dest}")
            
            # 复制处理报告
            report_source = os.path.join(self.output_dir, f"{outlet_id}_processing_report.json")
            if os.path.exists(report_source):
                report_dest = os.path.join(final_dir, f"{outlet_id}_processing_report.json")
                shutil.copy2(report_source, report_dest)
                final_output_files['processing_report'] = report_dest
                logger.info(f"已复制处理报告: {report_source} -> {report_dest}")
            
            logger.success(f"文件复制完成，共复制 {len(final_output_files)} 个文件到 {final_dir}")
            return final_output_files
            
        except Exception as e:
            logger.error(f"复制文件到最终输出目录失败: {str(e)}")
            return output_files


def main():
    """主函数，演示标准化流域划分功能"""

    logger.info("=== 标准化流域划分工具演示 ===")

    # 创建标准化流域划分器
    delineator = StandardizedDelineator()

    # 执行流域划分
    # results = delineator.delineate_watershed(
    #     lat=64.71072,
    #     lng=-21.60337,
    #     outlet_name="Nordhura River at Stekkur",
    #     outlet_id="6401070",
    #     area_estimate=507,
    #     high_res=True
    # )

    results = delineator.delineate_watershed(
        lat=46.279194,
        lng=130.673278,
        outlet_name="11106800",
        outlet_id="11106800",
        area_estimate=899.973,
        high_res=True
    )

    if results['validation_result']['is_valid']:
        logger.success("标准化流域划分成功完成！")
        logger.info("输出文件:")
        for file_type, file_path in results['output_files'].items():
            logger.info(f"  {file_type}: {file_path}")
    else:
        logger.error("标准化流域划分存在问题")
        logger.error(f"无效文件: {results['validation_result']['invalid_files']}")
        logger.error(f"缺失文件: {results['validation_result']['missing_files']}")


if __name__ == "__main__":
    main()